import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models

class Style(nn.Module):
    def __init__(self, pool_indexs='1+2+3'):
        super(Style, self).__init__()
        vgg_features = models.vgg16(pretrained=True).feature
        modules = [m for m in vgg_features]
        self.pool_indexs = pool_indexs.split('+')
        self.pool_indexs.sort()
        for p in self.pool_indexs:
            if p == '1':
                self.vgg = nn.ModuleList(nn.Sequential(modules[:5]))
            elif p == '2':
                if '1' in self.pool_indexs:
                    self.vgg.append(nn.Sequential(modules[5:10]))
                else:
                    self.vgg = nn.ModuleList(nn.Sequential(modules[:10]))
            else:
                if '2' in self.pool_indexs:
                    self.vgg.append(nn.Sequential(modules[9:19]))
                else:
                    self.vgg = nn.ModuleList(nn.Sequential(modules[:19]))
        
        self.vgg.requires_grad = False

    def forward(self, outs, gts):
        loss = 0
        for p_idx, _ in enumerate(self.pool_indexs):
            _, c, h, w = outs.size()
            outs = self.vgg[p_idx](outs)
            outs = outs.transpose() * outs
            with torch.no_grad():
                gts = self.vgg[p_idx](gts.detach())
                gts = gts.transpose() * gts
            k_idx = h * w * c
            loss += F.l1_loss(outs / k_idx, gts / k_idx)
        return loss